package org.korosoft.rusalad.syncnrun.tools;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.io.IOUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import java.io.*;
import java.math.BigInteger;
import java.security.*;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.*;

/**
 * Suncnrun key store implementation. Key store is similar to .ssh directory layout.
 *
 * @author Dmitry Korotkov
 * @since 1.0
 */
public class KeyStore {
    private static final Log log = LogFactory.getLog(KeyStore.class);
    public static final byte[] SENTINEL = new byte[]{'s', 's', 'h', '-', 'r', 's', 'a'};
    private final File path;

    public KeyStore(File path) {
        this.path = path;
        if (!path.exists()) {
            if (!path.mkdirs()) {
                throw new RuntimeException("Failed to create keystore directory: " + path.getAbsolutePath());
            }
        }
    }

    public PrivateKey getPrivateKey() {
        try {
            File keyFile = new File(path, "id_rsa");
            return loadPrivateKey(keyFile);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public PrivateKey loadPrivateKey(File keyFile) throws IOException, InvalidKeySpecException, NoSuchAlgorithmException {
        if (!keyFile.exists()) {
            log.warn("Private key is missing. generating a new one.");
            generatePrivateKey();
        }
        FileInputStream input = new FileInputStream(keyFile);
        List<String> strings = IOUtils.readLines(input);
        input.close();
        StringBuilder builder = new StringBuilder();
        for (String s : strings) {
            if (!(s.startsWith("-"))) {
                builder.append(s);
            }
        }
        byte[] keyData = Base64.decodeBase64(builder.toString());
        log.debug(String.format("Using public key from file %s", keyFile.getAbsolutePath()));
        return KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(keyData));
    }

    public void setPrivateKey(PrivateKey secretKey) {
        try {
            File keyFile = new File(path, "id_rsa");
            log.info("Writing private key to file " + keyFile);
            FileOutputStream output = new FileOutputStream(keyFile);
            output.write("-----BEGIN RSA PRIVATE KEY-----\n".getBytes());
            output.write(Base64.encodeBase64(secretKey.getEncoded(), true));
            output.write("-----END RSA PRIVATE KEY-----\b".getBytes());
            output.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public Collection<PublicKey> getPublicKeys() {
        File keyFile = new File(path, "authorized_keys");
        return loadPublicKeys(keyFile);
    }

    public static Collection<PublicKey> loadPublicKeys(File keyFile) {
        try {
            if (!keyFile.exists()) {
                return Collections.emptyList();
            }
            List<PublicKey> result = new ArrayList<PublicKey>();
            FileInputStream input = new FileInputStream(keyFile);
            List<String> strings = IOUtils.readLines(input);
            input.close();
            for (String s : strings) {
                String[] tokens = s.split(" +");
                if (tokens.length < 2) {
                    continue;
                }
                ByteArrayInputStream stream = new ByteArrayInputStream(Base64.decodeBase64(tokens[1]));
                final byte[] sentinel = read4251(stream);
                if (!Arrays.equals(sentinel, SENTINEL)) {
                    log.warn(String.format("Unexpected line in %s file: %s", keyFile.getAbsolutePath(), s));
                    continue;
                }
                BigInteger publicExponent = new BigInteger(read4251(stream));
                BigInteger modulus = new BigInteger(read4251(stream));
                RSAPublicKeySpec publicKeySpec = new RSAPublicKeySpec(modulus, publicExponent);
                result.add(KeyFactory.getInstance("RSA").generatePublic(publicKeySpec));
            }
            log.debug(String.format("Using %d public keys from file %s", result.size(), keyFile.getAbsolutePath()));
            return result;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public void addPublicKey(PublicKey publicKey) {
        File keyFile = new File(path, "authorized_keys");
        addPublicKey(publicKey, keyFile);
    }

    public void addPublicKey(PublicKey publicKey, File keyFile) {
        try {
            if (!(publicKey instanceof RSAPublicKey)) {
                throw new IllegalArgumentException("RSA Key expected but got " + publicKey.getAlgorithm());
            }
            BigInteger publicExponent = ((RSAPublicKey) publicKey).getPublicExponent();
            BigInteger modulus = ((RSAPublicKey) publicKey).getModulus();
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            write4251(baos, SENTINEL);
            write4251(baos, publicExponent.toByteArray());
            write4251(baos, modulus.toByteArray());
            baos.close();
            byte[] keyData = baos.toByteArray();
            OutputStream stream = new FileOutputStream(keyFile, true);
            stream.write("ssh-rsa ".getBytes());
            stream.write(Base64.encodeBase64(keyData, false));
            stream.write("\n".getBytes());
            stream.close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void generatePrivateKey() {
        KeyPairGenerator generator;
        try {
            generator = KeyPairGenerator.getInstance("RSA");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
        generator.initialize(1024);
        KeyPair keyPair = generator.genKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();
        setPrivateKey(privateKey);
        File publicKeyFile = new File(path, "id_rsa.pub");
        log.info("Generating private key to file " + publicKeyFile);
        if (publicKeyFile.exists()) {
            if (!publicKeyFile.delete()) {
                log.warn("Failed to delete " + publicKeyFile.getAbsolutePath());
            }
        }
        addPublicKey(publicKey, publicKeyFile);
    }

    private static byte[] read4251(InputStream stream) throws IOException {
        final int l1 = stream.read();
        final int l2 = stream.read();
        final int l3 = stream.read();
        final int l4 = stream.read();
        if (l1 < 0 || l2 < 0 || l3 < 0 || l4 < 0) {
            throw new IOException("Unexpected end of data");
        }
        final int len = (l1 << 24) + (l2 << 16) + (l3 << 8) + l4;
        if (len > 65536) {
            throw new IOException("Refusing to read object larger that 64K");
        }
        return IOUtils.toByteArray(stream, len);
    }

    private static void write4251(OutputStream outputStream, byte[] data) throws IOException {
        final int l = data.length;
        final int l1 = (l >> 24) & 255;
        final int l2 = (l >> 16) & 255;
        final int l3 = (l >> 8) & 255;
        final int l4 = l & 255;
        outputStream.write(l1);
        outputStream.write(l2);
        outputStream.write(l3);
        outputStream.write(l4);
        outputStream.write(data);
    }
}
